from typing import List, Dict, Any, Optional, Tuple
from app.graph_rag.service import graph_rag_service
import time
from functools import lru_cache
from ..core.logger import get_logger

logger = get_logger('graph_rag.mcp_tools')

class MCPGraphRAGTools:
    """MCP GraphRAG工具集"""
    
    def __init__(self):
        """初始化MCP工具集"""
        self.graph_rag_service = graph_rag_service
    
    def query_financial_data(self, query: str, timeout: int = 30, retry: int = 3) -> str:
        """
        查询金融数据，增加超时和重试机制
        
        Args:
            query: 查询语句
            timeout: 超时时间（秒）
            retry: 重试次数
            
        Returns:
            查询结果
        """
        for attempt in range(retry):
            try:
                logger.info(f"MCP工具查询金融数据: {query} (尝试次数: {attempt + 1})")
                start_time = time.time()
                
                # 使用缓存避免重复相同查询
                @lru_cache(maxsize=128)
                def cached_query(q):
                    return self.graph_rag_service.query(q)
                
                result = cached_query(query)
                
                elapsed_time = time.time() - start_time
                logger.info(f"查询耗时: {elapsed_time:.2f}秒")
                
                # 如果查询耗时过长，清除缓存重新执行
                if elapsed_time > timeout:
                    cached_query.cache_clear()
                    logger.warning(f"查询超时，清除缓存后重试")
                    continue
                    
                return result
                
            except Exception as e:
                logger.error(f"MCP工具查询金融数据时出错: {e}")
                if attempt < retry - 1:
                    logger.info(f"重试查询...")
                else:
                    return f"查询错误: {str(e)}"
        
        return f"查询超时: 超过最大重试次数"
    
    def analyze_fund_investment_pattern(self, fund_code: str, detail_level: str = "normal") -> str:
        """
        分析基金投资模式，增加细节控制参数
        
        Args:
            fund_code: 基金代码
            detail_level: 分析细节级别 ("normal" 或 "detailed")
            
        Returns:
            分析结果
        """
        try:
            logger.info(f"MCP工具分析基金投资模式: {fund_code} (细节级别: {detail_level})")
            
            # 根据不同细节级别采用不同的分析方法
            if detail_level == "detailed":
                result = self.graph_rag_service.detailed_analyze_investment_pattern(fund_code)
            else:
                result = self.graph_rag_service.analyze_investment_pattern(fund_code)
                
            return result
        except Exception as e:
            logger.error(f"MCP工具分析基金投资模式时出错: {e}")
            return f"分析错误: {str(e)}"
    
    def get_entity_relationships(self, entity_id: str) -> Dict[str, Any]:
        """
        获取实体关系
        
        Args:
            entity_id: 实体ID
            
        Returns:
            关系信息
        """
        try:
            logger.info(f"MCP工具获取实体关系: {entity_id}")
            result = self.graph_rag_service.get_entity_relationships(entity_id)
            return result
        except Exception as e:
            logger.error(f"MCP工具获取实体关系时出错: {e}")
            return {"error": f"获取关系错误: {str(e)}"}
    
    def find_entity_paths(self, source_id: str, target_id: str) -> List[Any]:
        """
        查找实体路径
        
        Args:
            source_id: 源实体ID
            target_id: 目标实体ID
            
        Returns:
            路径列表
        """
        try:
            logger.info(f"MCP工具查找实体路径: {source_id} -> {target_id}")
            result = self.graph_rag_service.find_paths(source_id, target_id)
            return result
        except Exception as e:
            logger.error(f"MCP工具查找实体路径时出错: {e}")
            return []

# 全局MCP工具实例
mcp_tools = MCPGraphRAGTools()

# 兼容旧接口的函数
def mcp_query_financial_data(query: str) -> str:
    """MCP兼容的金融数据查询函数"""
    return mcp_tools.query_financial_data(query)

def mcp_analyze_fund_investment_pattern(fund_code: str) -> str:
    """MCP兼容的基金投资模式分析函数"""
    return mcp_tools.analyze_fund_investment_pattern(fund_code)

def mcp_get_entity_relationships(entity_id: str) -> Dict[str, Any]:
    """MCP兼容的实体关系获取函数"""
    return mcp_tools.get_entity_relationships(entity_id)

def mcp_find_entity_paths(source_id: str, target_id: str) -> List[Any]:
    """MCP兼容的实体路径查找函数"""
    return mcp_tools.find_entity_paths(source_id, target_id)